import argparse
from collections import defaultdict
from model.net import CLSNet
from data.tvm_dataset import TVMDataSet
import torch
import os
from torch.utils.data import DataLoader
import tqdm
from loguru import logger

label_list = [
    'LSTM', 'UpSampling2D', 'UpSampling3D', 'Cropping2D', 'ZeroPadding2D', 'ZeroPadding3D', 'SeparableConv2D', 'Conv2D',
    'DepthwiseConv2D', 'ReLU', 'ThresholdedReLU', 'LeakyReLU', 'Softmax', 'ELU', 'Conv3D', 'Dense', 'Reshape',
    'Flatten', 'Concatenate', 'Average', 'Maximum', 'Minimum', 'Add', 'Subtract', 'Multiply', 'Dot', 'MaxPooling2D',
    'MaxPooling3D', 'AveragePooling2D', 'AveragePooling3D', 'GlobalMaxPooling2D', 'GlobalMaxPooling3D',
    'GlobalAveragePooling2D', 'GlobalAveragePooling3D', 'BatchNormalization'
]


def run(args):

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Prepare Data
    test_dataset = TVMDataSet(args.test_path, train_ratio=0.7, mode='test', w2v_model=args.w2v_model)
    test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=True, pin_memory=True)

    # Model Init
    net = CLSNet(in_d=200, out_classes=len(label_list), hiden=args.hsize, num_layers=args.num_layers)
    if not args.ckpt == "None":
        net.load_state_dict(torch.load(args.ckpt))
    net.to(device)

    # Test
    logger.info("Start evaluation ...")
    net.eval()
    tmp_count = 0
    succ_count = 0
    pbar = tqdm.tqdm(total=len(test_dataset))
    succ_dict = defaultdict(lambda: 0)
    total_dict = defaultdict(lambda: 0)
    for input, target in test_dataloader:
        pbar.update(1)
        pred = net(input, device)
        if pred is None:
            continue
        pred = pred[0].topk(1)[1][0]
        if pred == target[0].to(device):
            succ_count += 1
            succ_dict[int(target[0][0])] += 1
        tmp_count += 1
        total_dict[int(target[0][0])] += 1
    pbar.close()
    logger.info("Accuracy: {}".format(succ_count / tmp_count))
    if args.show_failure_case:
        for k in total_dict:
            logger.info("Target:[{}] Total:{}, Succ:{}".format(label_list[k], total_dict[k], succ_dict[k]))
    logger.info("Finished!")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='TVM Reversion')
    parser.add_argument("--hsize", default=200, type=int)
    parser.add_argument("--num_layers", default=2, type=int)
    parser.add_argument("--ckpt", default="None", type=str)
    parser.add_argument("--test_path", type=str)
    parser.add_argument("--gpu", default="0", type=str)
    parser.add_argument("--w2v_model", default="w2v-new.model", type=str)
    parser.add_argument("-sfc", "--show_failure_case", dest="show_failure_case", action="store_true")
    parser.set_defaults(show_failure_case=False)
    args = parser.parse_args()
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    run(args)
